#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
FlinkCEP 示例数据生成器
用于生成测试数据并发送到Kafka主题

使用方法:
python sample-data-generator.py --kafka-server localhost:9092 --topic input-events
"""

import json
import time
import random
import argparse
from datetime import datetime
from kafka import KafkaProducer
import uuid

class SampleDataGenerator:
    def __init__(self, kafka_server, topic):
        self.producer = KafkaProducer(
            bootstrap_servers=[kafka_server],
            value_serializer=lambda v: json.dumps(v, ensure_ascii=False).encode('utf-8')
        )
        self.topic = topic
        self.user_ids = [f"user{i:03d}" for i in range(1, 21)]  # 20个用户
        self.sources = ["AUTH_SERVICE", "API_GATEWAY", "DATABASE", "CACHE_SERVICE", "PAYMENT_SERVICE"]
        self.event_types = [
            "LOGIN_SUCCESS", "LOGIN_FAILURE", "API_CALL", "DATABASE_QUERY", 
            "PAYMENT_PROCESS", "DATA_ANOMALY", "PERFORMANCE_WARNING", 
            "SYSTEM_ERROR", "SECURITY_ALERT", "UNAUTHORIZED_ACCESS"
        ]
        self.severities = ["INFO", "WARNING", "ERROR", "CRITICAL"]
        
    def generate_normal_event(self):
        """生成正常事件"""
        return {
            "id": str(uuid.uuid4()),
            "type": random.choice(["LOGIN_SUCCESS", "API_CALL", "DATABASE_QUERY"]),
            "source": random.choice(self.sources),
            "timestamp": int(time.time() * 1000),
            "data": "正常操作",
            "severity": "INFO",
            "userId": random.choice(self.user_ids)
        }
    
    def generate_error_event(self, user_id=None):
        """生成错误事件"""
        return {
            "id": str(uuid.uuid4()),
            "type": random.choice(["SYSTEM_ERROR", "DATABASE_ERROR", "API_ERROR"]),
            "source": random.choice(self.sources),
            "timestamp": int(time.time() * 1000),
            "data": "系统发生错误",
            "severity": "ERROR",
            "userId": user_id or random.choice(self.user_ids)
        }
    
    def generate_login_failure_event(self, user_id=None):
        """生成登录失败事件"""
        return {
            "id": str(uuid.uuid4()),
            "type": "LOGIN_FAILURE",
            "source": "AUTH_SERVICE",
            "timestamp": int(time.time() * 1000),
            "data": "用户登录失败",
            "severity": "ERROR",
            "userId": user_id or random.choice(self.user_ids)
        }
    
    def generate_security_event(self, user_id=None):
        """生成安全事件"""
        return {
            "id": str(uuid.uuid4()),
            "type": random.choice(["SECURITY_ALERT", "UNAUTHORIZED_ACCESS", "INTRUSION_ATTEMPT"]),
            "source": "SECURITY_SERVICE",
            "timestamp": int(time.time() * 1000),
            "data": "检测到安全威胁",
            "severity": random.choice(["ERROR", "CRITICAL"]),
            "userId": user_id or random.choice(self.user_ids)
        }
    
    def generate_data_anomaly_event(self, source=None):
        """生成数据异常事件"""
        return {
            "id": str(uuid.uuid4()),
            "type": "DATA_ANOMALY",
            "source": source or random.choice(self.sources),
            "timestamp": int(time.time() * 1000),
            "data": "数据异常检测",
            "severity": "HIGH",
            "userId": random.choice(self.user_ids)
        }
    
    def generate_performance_warning_event(self):
        """生成性能警告事件"""
        return {
            "id": str(uuid.uuid4()),
            "type": "PERFORMANCE_WARNING",
            "source": random.choice(self.sources),
            "timestamp": int(time.time() * 1000),
            "data": "系统性能警告",
            "severity": "WARNING",
            "userId": random.choice(self.user_ids)
        }
    
    def send_event(self, event):
        """发送事件到Kafka"""
        try:
            self.producer.send(self.topic, event)
            print(f"发送事件: {event['type']} - {event['severity']} - 用户: {event['userId']}")
        except Exception as e:
            print(f"发送事件失败: {e}")
    
    def simulate_high_frequency_errors(self, user_id, count=5):
        """模拟高频错误场景"""
        print(f"\n=== 模拟高频错误场景 (用户: {user_id}) ===")
        for i in range(count):
            event = self.generate_error_event(user_id)
            self.send_event(event)
            time.sleep(2)  # 2秒间隔
    
    def simulate_abnormal_login(self, user_id, count=6):
        """模拟异常登录场景"""
        print(f"\n=== 模拟异常登录场景 (用户: {user_id}) ===")
        for i in range(count):
            event = self.generate_login_failure_event(user_id)
            self.send_event(event)
            time.sleep(10)  # 10秒间隔
    
    def simulate_security_threat(self, user_id, count=3):
        """模拟安全威胁场景"""
        print(f"\n=== 模拟安全威胁场景 (用户: {user_id}) ===")
        for i in range(count):
            event = self.generate_security_event(user_id)
            self.send_event(event)
            time.sleep(5)  # 5秒间隔
    
    def simulate_data_anomaly(self, source, count=4):
        """模拟数据异常场景"""
        print(f"\n=== 模拟数据异常场景 (来源: {source}) ===")
        for i in range(count):
            event = self.generate_data_anomaly_event(source)
            self.send_event(event)
            time.sleep(2)  # 2秒间隔
    
    def simulate_performance_degradation(self):
        """模拟性能降级场景"""
        print(f"\n=== 模拟性能降级场景 ===")
        # 先发送性能警告
        warning_event = self.generate_performance_warning_event()
        self.send_event(warning_event)
        
        time.sleep(30)  # 等待30秒
        
        # 然后发送系统错误
        error_event = {
            "id": str(uuid.uuid4()),
            "type": "SYSTEM_ERROR",
            "source": warning_event["source"],
            "timestamp": int(time.time() * 1000),
            "data": "系统错误",
            "severity": "ERROR",
            "userId": warning_event["userId"]
        }
        self.send_event(error_event)
    
    def generate_continuous_data(self, duration_minutes=10):
        """生成连续的测试数据"""
        print(f"\n=== 开始生成连续数据 (持续 {duration_minutes} 分钟) ===")
        end_time = time.time() + (duration_minutes * 60)
        
        while time.time() < end_time:
            # 90% 正常事件，10% 异常事件
            if random.random() < 0.9:
                event = self.generate_normal_event()
            else:
                event_type = random.choice(["error", "login_failure", "security"])
                if event_type == "error":
                    event = self.generate_error_event()
                elif event_type == "login_failure":
                    event = self.generate_login_failure_event()
                else:
                    event = self.generate_security_event()
            
            self.send_event(event)
            time.sleep(random.uniform(0.5, 3.0))  # 随机间隔
    
    def close(self):
        """关闭生产者"""
        self.producer.close()

def main():
    parser = argparse.ArgumentParser(description='FlinkCEP 示例数据生成器')
    parser.add_argument('--kafka-server', default='localhost:9092', help='Kafka服务器地址')
    parser.add_argument('--topic', default='input-events', help='Kafka主题名称')
    parser.add_argument('--scenario', choices=['all', 'high-freq-error', 'abnormal-login', 
                                              'security-threat', 'data-anomaly', 
                                              'performance-degradation', 'continuous'], 
                       default='all', help='测试场景')
    parser.add_argument('--duration', type=int, default=10, help='连续数据生成持续时间(分钟)')
    
    args = parser.parse_args()
    
    print(f"连接到Kafka: {args.kafka_server}, 主题: {args.topic}")
    
    generator = SampleDataGenerator(args.kafka_server, args.topic)
    
    try:
        if args.scenario == 'all' or args.scenario == 'high-freq-error':
            generator.simulate_high_frequency_errors("user001")
            time.sleep(5)
        
        if args.scenario == 'all' or args.scenario == 'abnormal-login':
            generator.simulate_abnormal_login("user002")
            time.sleep(5)
        
        if args.scenario == 'all' or args.scenario == 'security-threat':
            generator.simulate_security_threat("user003")
            time.sleep(5)
        
        if args.scenario == 'all' or args.scenario == 'data-anomaly':
            generator.simulate_data_anomaly("DATABASE")
            time.sleep(5)
        
        if args.scenario == 'all' or args.scenario == 'performance-degradation':
            generator.simulate_performance_degradation()
            time.sleep(5)
        
        if args.scenario == 'continuous':
            generator.generate_continuous_data(args.duration)
        
        print("\n数据生成完成!")
        
    except KeyboardInterrupt:
        print("\n用户中断，停止数据生成")
    except Exception as e:
        print(f"\n发生错误: {e}")
    finally:
        generator.close()

if __name__ == '__main__':
    main()